"""
比较运算，eq, equal, ge/gt, le/lt, ne
排序：sort,topk
有界无界：isfinite, isinf, isnan
"""

import torch
import numpy as np


def test_compare():
    """比较时，维度需要一致"""
    a = torch.rand(2, 3)
    b = torch.rand(2, 3)
    print("a")
    print(a)
    print("b")
    print(b)

    print("========= eq(逐个元素比较, 返回Tensor(元素为True/False)) ========")
    print(torch.eq(a, b))

    print("========= equal(两个tensor整体比较（元素和shape都相同）, 返回True/False) ========")
    print(torch.equal(a, b))

    print("========= ge / gt  (逐个元素比较, 返回Tensor(元素为True/False)) ========")
    print("ge")
    print(torch.ge(a, b))
    print("gt")
    print(torch.gt(a, b))

    print("========= le / lt  (逐个元素比较, 返回Tensor(元素为True/False)) ========")
    print("le")
    print(torch.le(a, b))
    print("lt")
    print(torch.lt(a, b))

    print("========= ne  (逐个元素比较, 返回Tensor(元素为True/False)) ========")
    print("ne")
    print(torch.ne(a, b))


def test_sort():
    a = torch.tensor([[1, 4, 4, 3, 5],
                      [2, 3, 1, 3, 5]])
    print(a)
    print(a.shape)

    print("===== sort(a) =====")
    # dim=0 按列，dim=1 按行，高维依次在各自维度
    print(torch.sort(a, descending=True, dim=1))


def test_topk():
    print("===== TOPK 从大到小 ====")
    a = torch.tensor([[2, 4, 3, 1, 5],
                      [2, 3, 5, 1, 4]])
    print("===== a ====")
    print(a)
    print(a.shape)
    print()

    print("""===== torch.topk(a, k=1, dim=0) ====""")
    # dim=0 按列，dim=1 按行，高维依次在各自维度
    print(torch.topk(a, k=1, dim=0))
    print()

    print("""===== torch.topk(a, k=2, dim=1) ====""")
    # dim=0 按列，dim=1 按行，高维依次在各自维度
    print(torch.topk(a, k=2, dim=1))
    print()


def test_kthvalue():
    print("===== kthvalue 取k小的数 ====")
    a = torch.tensor([[2, 4, 3, 1, 5],
                      [2, 3, 5, 1, 4]])
    print("===== a ====")
    print(a)
    print(a.shape)
    print()

    print("""===== torch.kthvalue(a, k=2, dim=0) ====""")
    # dim=0 按列，dim=1 按行，高维依次在各自维度
    print(torch.kthvalue(a, k=2, dim=0))
    print()

    print("""===== torch.kthvalue(a, k=2, dim=1) ====""")
    # dim=0 按列，dim=1 按行，高维依次在各自维度
    print(torch.kthvalue(a, k=2, dim=1))
    print()


def test_isfinite_isinf_isnan():
    print("""===== isinfite ====""")
    a = torch.rand(2, 3)
    print("===== a")
    print(a)

    print("===== torch.isfinite(a)  是否有界 ====")
    print(torch.isfinite(a))
    print()

    print("===== a/0 ====")
    print(a / 0)
    print()

    print("===== torch.isfinite(a/0)  是否有界 ====")
    print(torch.isfinite(a / 0))
    print()

    print("===== torch.isinf(a/0)  是否无界 ====")
    print(torch.isinf(a / 0))
    print()

    print("===== torch.isnan(a/0)  是否nan ====")
    print(torch.isnan(a / 0))
    print()

    b = torch.tensor([1, 2, np.nan])
    print("===== b ====")
    print(b)
    print("===== torch.isnan(b)  是否nan ====")
    print(torch.isnan(b))
    print()


if __name__ == '__main__':
    # test_compare()
    # test_sort()
    # test_topk()
    # test_kthvalue()
    test_isfinite_isinf_isnan()
